{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -c \"import monai\" || pip install -q \"monai-weekly[gdown, nibabel, tqdm]\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gdown\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import jit\n",
    "\n",
    "from monai.apps.deepgrow.transforms import (\n",
    "    AddGuidanceFromPointsd,\n",
    "    AddGuidanceSignald,\n",
    "    ResizeGuidanced,\n",
    "    RestoreLabeld,\n",
    "    SpatialCropGuidanced,\n",
    ")\n",
    "from monai.transforms import (\n",
    "    AsChannelFirstd,\n",
    "    Spacingd,\n",
    "    LoadImaged,\n",
    "    AddChanneld,\n",
    "    NormalizeIntensityd,\n",
    "    EnsureTyped,\n",
    "    ToNumpyd,\n",
    "    Activationsd,\n",
    "    AsDiscreted,\n",
    "    Resized\n",
    ")\n",
    "\n",
    "max_epochs = 1\n",
    "\n",
    "\n",
    "def draw_points(guidance, slice_idx):\n",
    "    if guidance is None:\n",
    "        return\n",
    "    colors = ['r+', 'b+']\n",
    "    for color, points in zip(colors, guidance):\n",
    "        for p in points:\n",
    "            if p[0] != slice_idx:\n",
    "                continue\n",
    "            p1 = p[-1]\n",
    "            p2 = p[-2]\n",
    "            plt.plot(p1, p2, color, 'MarkerSize', 30)\n",
    "\n",
    "\n",
    "def show_image(image, label, guidance=None, slice_idx=None):\n",
    "    plt.figure(\"check\", (12, 6))\n",
    "    plt.subplot(1, 2, 1)\n",
    "    plt.title(\"image\")\n",
    "    plt.imshow(image, cmap=\"gray\")\n",
    "\n",
    "    if label is not None:\n",
    "        masked = np.ma.masked_where(label == 0, label)\n",
    "        plt.imshow(masked, 'jet', interpolation='none', alpha=0.7)\n",
    "\n",
    "    draw_points(guidance, slice_idx)\n",
    "    plt.colorbar()\n",
    "\n",
    "    if label is not None:\n",
    "        plt.subplot(1, 2, 2)\n",
    "        plt.title(\"label\")\n",
    "        plt.imshow(label)\n",
    "        plt.colorbar()\n",
    "        # draw_points(guidance, slice_idx)\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def print_data(data):\n",
    "    for k in data:\n",
    "        v = data[k]\n",
    "\n",
    "        d = type(v)\n",
    "        if type(v) in (int, float, bool, str, dict, tuple):\n",
    "            d = v\n",
    "        elif hasattr(v, 'shape'):\n",
    "            d = v.shape\n",
    "\n",
    "        if k in ('image_meta_dict', 'label_meta_dict'):\n",
    "            for m in data[k]:\n",
    "                print('{} Meta:: {} => {}'.format(k, m, data[k][m]))\n",
    "        else:\n",
    "            print('Data key: {} = {}'.format(k, d))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download data and model\n",
    "\n",
    "resource = \"https://drive.google.com/uc?id=1cIlDXWx4pEFpldoIXMEe-5JeaOxzB05Z\"\n",
    "dst = \"_image.nii.gz\"\n",
    "\n",
    "if not os.path.exists(dst):\n",
    "    gdown.download(resource, dst, quiet=False)\n",
    "\n",
    "resource = \"https://drive.google.com/uc?id=1BcU4Z-wdkw7xjydDNd28iVBUVDJYKqCO\"\n",
    "dst = \"deepgrow_3d.ts\"\n",
    "if not os.path.exists(dst):\n",
    "    gdown.download(resource, dst, quiet=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pre Processing\n",
    "roi_size = [256, 256]\n",
    "model_size = [128, 192, 192]\n",
    "pixdim = (1.0, 1.0, 1.0)\n",
    "dimensions = 3\n",
    "\n",
    "data = {\n",
    "    'image': '_image.nii.gz',\n",
    "    'foreground': [[66, 180, 105], [66, 180, 145]],\n",
    "    'background': [],\n",
    "}\n",
    "slice_idx = original_slice_idx = data['foreground'][0][2]\n",
    "\n",
    "pre_transforms = [\n",
    "    LoadImaged(keys='image'),\n",
    "    AsChannelFirstd(keys='image'),\n",
    "    Spacingd(keys='image', pixdim=pixdim, mode='bilinear'),\n",
    "    AddGuidanceFromPointsd(ref_image='image', guidance='guidance', foreground='foreground', background='background',\n",
    "                           dimensions=dimensions),\n",
    "    AddChanneld(keys='image'),\n",
    "    SpatialCropGuidanced(keys='image', guidance='guidance', spatial_size=roi_size),\n",
    "    Resized(keys='image', spatial_size=model_size, mode='area'),\n",
    "    ResizeGuidanced(guidance='guidance', ref_image='image'),\n",
    "    NormalizeIntensityd(keys='image', subtrahend=208.0, divisor=388.0),\n",
    "    AddGuidanceSignald(image='image', guidance='guidance'),\n",
    "    EnsureTyped(keys='image')\n",
    "]\n",
    "\n",
    "original_image = None\n",
    "for t in pre_transforms:\n",
    "    tname = type(t).__name__\n",
    "    data = t(data)\n",
    "    image = data['image']\n",
    "    label = data.get('label')\n",
    "    guidance = data.get('guidance')\n",
    "\n",
    "    print(\"{} => image shape: {}\".format(tname, image.shape))\n",
    "\n",
    "    guidance = guidance if guidance else [np.roll(data['foreground'], 1).tolist(), []]\n",
    "    slice_idx = guidance[0][0][0] if guidance else slice_idx\n",
    "    print('Guidance: {}; Slice Idx: {}'.format(guidance, slice_idx))\n",
    "    if tname == 'Resized':\n",
    "        continue\n",
    "\n",
    "    image = image[:, :, slice_idx] if tname in ('LoadImaged') else image[slice_idx] if tname in (\n",
    "        'AsChannelFirstd', 'Spacingd', 'AddGuidanceFromPointsd') else image[0][slice_idx]\n",
    "    label = None\n",
    "\n",
    "    show_image(image, label, guidance, slice_idx)\n",
    "    if tname == 'LoadImaged':\n",
    "        original_image = data['image']\n",
    "    if tname == 'AddChanneld':\n",
    "        original_image_slice = data['image']\n",
    "    if tname == 'SpatialCropGuidanced':\n",
    "        spatial_image = data['image']\n",
    "\n",
    "image = data['image']\n",
    "label = data.get('label')\n",
    "guidance = data.get('guidance')\n",
    "for i in range(image.shape[1]):\n",
    "    print('Slice Idx: {}'.format(i))\n",
    "    # show_image(image[0][i], None, guidance, i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluation\n",
    "model_path = 'deepgrow_3d.ts'\n",
    "model = jit.load(model_path)\n",
    "model.cuda()\n",
    "model.eval()\n",
    "\n",
    "inputs = data['image'][None].cuda()\n",
    "with torch.no_grad():\n",
    "    outputs = model(inputs)\n",
    "outputs = outputs[0]\n",
    "data['pred'] = outputs\n",
    "\n",
    "post_transforms = [\n",
    "    Activationsd(keys='pred', sigmoid=True),\n",
    "    AsDiscreted(keys='pred', threshold_values=True, logit_thresh=0.5),\n",
    "    ToNumpyd(keys='pred'),\n",
    "    RestoreLabeld(keys='pred', ref_image='image', mode='nearest'),\n",
    "]\n",
    "\n",
    "pred = None\n",
    "for t in post_transforms:\n",
    "    tname = type(t).__name__\n",
    "\n",
    "    data = t(data)\n",
    "    image = data['image']\n",
    "    label = data['pred']\n",
    "    print(\"{} => image shape: {}, pred shape: {}; slice_idx: {}\".format(tname, image.shape, label.shape, slice_idx))\n",
    "\n",
    "    if tname in 'RestoreLabeld':\n",
    "        pred = label\n",
    "\n",
    "        image = original_image[:, :, original_slice_idx]\n",
    "        label = label[original_slice_idx]\n",
    "        print(\"PLOT:: {} => image shape: {}, pred shape: {}; min: {}, max: {}, sum: {}\".format(\n",
    "            tname, image.shape, label.shape, np.min(label), np.max(label), np.sum(label)))\n",
    "        show_image(image, label)\n",
    "    elif tname == 'xToNumpyd':\n",
    "        for i in range(label.shape[1]):\n",
    "            img = image[0, i, :, :].detach().cpu().numpy() if torch.is_tensor(image) else image[0][i]\n",
    "            lab = label[0, i, :, :].detach().cpu().numpy() if torch.is_tensor(label) else label[0][i]\n",
    "            if np.sum(lab) > 0:\n",
    "                print(\"PLOT:: {} => image shape: {}, pred shape: {}; min: {}, max: {}, sum: {}\".format(\n",
    "                    i, img.shape, lab.shape, np.min(lab), np.max(lab), np.sum(lab)))\n",
    "                show_image(img, lab)\n",
    "    else:\n",
    "        image = image[0, slice_idx, :, :].detach().cpu().numpy() if torch.is_tensor(image) else image[0][slice_idx]\n",
    "        label = label[0, slice_idx, :, :].detach().cpu().numpy() if torch.is_tensor(label) else label[0][slice_idx]\n",
    "        print(\"PLOT:: {} => image shape: {}, pred shape: {}; min: {}, max: {}, sum: {}\".format(\n",
    "            tname, image.shape, label.shape, np.min(label), np.max(label), np.sum(label)))\n",
    "        show_image(image, label)\n",
    "\n",
    "for i in range(pred.shape[0]):\n",
    "    image = original_image[:, :, i]\n",
    "    label = pred[i, :, :]\n",
    "    if np.sum(label) == 0:\n",
    "        continue\n",
    "\n",
    "    print(\"Final PLOT:: {} => image shape: {}, pred shape: {}; min: {}, max: {}, sum: {}\".format(\n",
    "        i, image.shape, label.shape, np.min(label), np.max(label), np.sum(label)))\n",
    "    show_image(image, label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "pred = data['pred']\n",
    "meta_data = data['pred_meta_dict']\n",
    "affine = meta_data.get(\"affine\", None)\n",
    "\n",
    "pred = np.moveaxis(pred, 0, -1)\n",
    "print('Prediction NII shape: {}'.format(pred.shape))\n",
    "\n",
    "# file_name = 'result_label.nii.gz'\n",
    "# write_nifti(pred, file_name=file_name)\n",
    "# print('Prediction saved at: {}'.format(file_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# remove downloaded files\n",
    "os.remove('_image.nii.gz')\n",
    "os.remove('deepgrow_3d.ts')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
